package utils

import (
	"encoding/json"
	"fmt"
	"geekai/core/types"
	logger2 "geekai/logger"
	"geekai/model"
	"github.com/imroc/req/v3"
	"github.com/pkoukk/tiktoken-go"
	"gorm.io/gorm"
	"io"
	"time"
)

var log = logger2.GetLogger()

// defaultModel 默认模型
var defaultModel = "glm-4-plus"

// OpenAIResponse 接收openai返回结果
type OpenAIResponse struct {
	Model   string `json:"model"`
	Choices []struct {
		Index   int `json:"index"` // 结果索引
		Message struct {
			Role    string `json:"role"`    // 角色
			Content string `json:"content"` // 消息内容
		} `json:"message"`
		FinishReason string `json:"finish_reason"` // 模型推理终止原因
	} `json:"choices"`
}

// OpenAIRequest 向大模型发送请求
func OpenAIRequest(db *gorm.DB, prompt string, modelId int) (string, error) {
	messages := make([]interface{}, 1)
	messages[0] = types.Message{
		Role:    "user",
		Content: prompt,
	}
	return SendOpenAIMessage(db, messages, modelId)
}

// SendOpenAIMessage 向大模型发送请求
func SendOpenAIMessage(db *gorm.DB, messages []interface{}, modelId int) (string, error) {
	var chatModel model.ChatModel
	// 根据模型Id，查询模型信息
	db.Where("id", modelId).First(&chatModel)
	// 判断模型值是否为空，如果为空就指定默认的模型
	if chatModel.Value == "" {
		chatModel.Value = defaultModel // 默认使用的模型
	}
	// 查询模型对应的 Api Key
	var apiKey model.ApiKey
	session := db.Session(&gorm.Session{}).Where("type", "chat").Where("enabled", true)
	if chatModel.KeyId > 0 {
		session = session.Where("id", chatModel.KeyId)
	}
	err := session.First(&apiKey).Error
	if err != nil {
		return "", fmt.Errorf("查询模型ApiKey出现错误：%v", err)
	}
	var response OpenAIResponse
	client := req.C()
	if len(apiKey.ProxyURL) > 5 {
		client.SetProxyURL(apiKey.ApiURL)
	}
	// 发送请求地址
	apiURL := fmt.Sprintf("%s/v1/chat/completions", apiKey.ApiURL)
	log.Infof("向大模型发送请求：%s，API KEY：%s，PROXY：%s，MODEL：%s", apiURL, apiKey.Value, apiKey.ProxyURL, chatModel.Value)
	r, err := client.R().SetHeader("Body-Type", "application/json").
		SetHeader("Authorization", "Bearer "+apiKey.Value).
		SetBody(types.ApiRequest{
			Model:       chatModel.Value,
			Temperature: 0.9,
			MaxTokens:   1024,
			Stream:      false,
			Messages:    messages,
		}).Post(apiURL)
	if err != nil {
		return "", fmt.Errorf("请求 OpenAI API失败：%v", err)
	}
	if r.IsErrorState() {
		return "", fmt.Errorf("请求 OpenAI API失败：%v", r.Status)
	}
	body, _ := io.ReadAll(r.Body)
	err = json.Unmarshal(body, &response)
	if err != nil {
		return "", fmt.Errorf("解析API数据失败：%v，%s", err, string(body))
	}
	// 更新API KEY 的最后使用时间
	db.Model(&apiKey).UpdateColumn("last_used_at", time.Now().Unix())
	return response.Choices[0].Message.Content, nil
}

// CalcTokens 计算文本在指定模型消耗的token数量
func CalcTokens(text string, model string) (int, error) {
	// 根据模型获取编码
	encoding, ok := tiktoken.MODEL_TO_ENCODING[model]
	if !ok {
		encoding = "cl100k_base"
	}
	tke, err := tiktoken.GetEncoding(encoding)
	if err != nil {
		return 0, fmt.Errorf("GetEncoding：%v", err)
	}
	// 将文本进行编码，得到token数量
	token := tke.Encode(text, nil, nil)
	return len(token), nil
}
